#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
基于Vosk实现中文语音识别
可通过键盘控制录音开始和结束
自动检测和加载模型vosk-model-small-cn-0.22
"""

import os
import sys
import json
import time
import threading
import queue

# 尝试导入必要的库
try:
    import pyaudio
except ImportError:
    print("错误：未找到PyAudio库。请安装：pip install pyaudio")
    sys.exit(1)

try:
    from vosk import Model, KaldiRecognizer, SetLogLevel
    # 设置日志级别为仅显示错误
    SetLogLevel(0)
except ImportError:
    print("错误：未找到Vosk库。请安装：pip install vosk==0.3.45")
    sys.exit(1)

# 检查是否可以使用msvcrt进行键盘输入检测（仅限Windows）
try:
    import msvcrt
    WINDOWS = True
except ImportError:
    WINDOWS = False
    print("提示：非Windows系统，键盘控制功能可能受限")

# 定义模型路径
MODEL_NAME = "vosk-model-small-cn-0.22"
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

# 支持两种常见的模型路径结构
MODEL_PATHS = [
    os.path.join(CURRENT_DIR, MODEL_NAME),                # 单层目录
    os.path.join(CURRENT_DIR, MODEL_NAME, MODEL_NAME)     # 双层目录
]

def check_model_path():
    """检查模型路径是否有效并返回可用的路径"""
    for path in MODEL_PATHS:
        print(f"检查模型路径: {path}")
        if os.path.exists(path):
            # 检查必要的子目录
            required = ['conf', 'am', 'graph']
            missing = [d for d in required if not os.path.exists(os.path.join(path, d))]
            
            if not missing:
                print(f"✓ 找到有效的模型路径: {path}")
                return path
            else:
                print(f"✗ 模型路径存在但缺少必要子目录: {', '.join(missing)}")
    
    print("错误: 无法找到有效的模型路径")
    print(f"请确保 {MODEL_NAME} 文件夹存在，并包含必要的子目录")
    sys.exit(1)

# 音频参数
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000
CHUNK = 4000

class VoiceSpeechRecognizer:
    """语音识别器类"""
    
    def __init__(self):
        self.audio = pyaudio.PyAudio()
        self.model_path = check_model_path()
        self.audio_queue = queue.Queue()
        self.recording = False
        self.results = []
        self.recognizer = None
        self.stream = None
        self.process_thread = None
        
        # 尝试加载模型
        try:
            self.model = Model(self.model_path)
            print("模型加载成功")
        except Exception as e:
            print(f"模型加载失败: {e}")
            try:
                # 尝试替代路径格式（解决某些Windows路径问题）
                alt_path = self.model_path.replace('\\', '/')
                self.model = Model(alt_path)
                print("使用替代路径格式成功加载模型")
            except Exception as e:
                print(f"所有模型加载尝试均失败: {e}")
                sys.exit(1)
    
    def list_audio_devices(self):
        """列出所有可用的音频输入设备"""
        print("\n可用的音频输入设备:")
        print("=" * 50)
        print(f"{'ID':<4}| {'设备名称':<40}")
        print("-" * 50)
        
        input_devices = []
        for i in range(self.audio.get_device_count()):
            device_info = self.audio.get_device_info_by_index(i)
            if device_info.get('maxInputChannels') > 0:
                print(f"{i:<4}| {device_info.get('name')}")
                input_devices.append(i)
        
        print("=" * 50)
        return input_devices
    
    def get_default_device(self):
        """获取默认输入设备ID"""
        try:
            info = self.audio.get_host_api_info_by_index(0)
            return info.get('defaultInputDevice')
        except:
            # 如果无法获取默认设备，尝试找到第一个输入设备
            for i in range(self.audio.get_device_count()):
                device_info = self.audio.get_device_info_by_index(i)
                if device_info.get('maxInputChannels') > 0:
                    return i
            return None
    
    def audio_callback(self, in_data, frame_count, time_info, status):
        """音频回调函数，将收到的音频数据放入队列"""
        if self.recording:
            self.audio_queue.put(in_data)
        return (in_data, pyaudio.paContinue)
    
    def process_audio_data(self):
        """处理音频数据并执行识别"""
        self.recognizer = KaldiRecognizer(self.model, RATE)
        
        while True:
            # 检查是否应退出线程
            if not self.recording and self.audio_queue.empty():
                break
            
            try:
                # 获取音频数据（设置超时防止无限阻塞）
                data = self.audio_queue.get(block=True, timeout=0.5)
                
                # 处理音频数据
                if self.recognizer.AcceptWaveform(data):
                    result = json.loads(self.recognizer.Result())
                    text = result.get('text', '').strip()
                    if text:
                        print(f"\n识别结果: {text}")
                        self.results.append(text)
                else:
                    # 显示部分结果
                    partial = json.loads(self.recognizer.PartialResult())
                    partial_text = partial.get('partial', '')
                    if partial_text:
                        print(f"\r当前识别: {partial_text}", end='', flush=True)
                
                self.audio_queue.task_done()
            
            except queue.Empty:
                pass  # 队列为空，继续等待
    
    def start_recording(self, device_id):
        """开始录音"""
        if self.recording:
            return
        
        # 清空之前的结果
        self.results.clear()
        self.audio_queue = queue.Queue()
        
        # 打开音频流
        try:
            self.stream = self.audio.open(
                format=FORMAT,
                channels=CHANNELS,
                rate=RATE,
                input=True,
                input_device_index=device_id,
                frames_per_buffer=CHUNK,
                stream_callback=self.audio_callback
            )
            
            self.recording = True
            
            # 启动处理线程
            self.process_thread = threading.Thread(target=self.process_audio_data)
            self.process_thread.daemon = True
            self.process_thread.start()
            
            print("\n✓ 开始录音...")
            print("按空格键暂停/继续录音，按Esc键退出")
            
        except Exception as e:
            print(f"\n错误: 无法打开音频设备 (ID: {device_id})")
            print(f"详细信息: {str(e)}")
            self.recording = False
    
    def stop_recording(self):
        """停止录音"""
        if not self.recording:
            return
        
        print("\n✓ 停止录音...")
        self.recording = False
        
        # 关闭音频流
        if self.stream is not None:
            self.stream.stop_stream()
            self.stream.close()
            self.stream = None
        
        # 等待处理线程完成
        if self.process_thread is not None:
            self.process_thread.join(timeout=2.0)
            self.process_thread = None
        
        # 处理最后的数据
        if self.recognizer is not None:
            final_result = json.loads(self.recognizer.FinalResult())
            final_text = final_result.get('text', '').strip()
            if final_text:
                print(f"最终识别: {final_text}")
                self.results.append(final_text)
        
        # 显示完整结果
        if self.results:
            print("\n===== 完整识别结果 =====")
            print("-" * 50)
            full_text = " ".join(self.results)
            print(full_text)
            print("-" * 50)
            
            # 保存结果到文件
            try:
                with open("识别结果.txt", "w", encoding="utf-8") as f:
                    f.write(f"语音识别结果 - {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                    f.write("-" * 50 + "\n")
                    f.write(full_text + "\n")
                print("结果已保存到: 识别结果.txt")
            except Exception as e:
                print(f"保存结果失败: {e}")
        else:
            print("没有识别到任何内容")
    
    def close(self):
        """关闭资源"""
        self.stop_recording()
        self.audio.terminate()

def main():
    """主函数"""
    print("\n===== 中文语音识别程序 =====")
    print("基于Vosk的离线语音识别")
    print("=" * 50)
    
    recognizer = VoiceSpeechRecognizer()
    
    # 列出可用设备
    input_devices = recognizer.list_audio_devices()
    if not input_devices:
        print("错误：未找到任何输入设备")
        return
    
    # 选择设备
    default_device = recognizer.get_default_device()
    if default_device is not None and default_device in input_devices:
        print(f"\n默认使用设备 [{default_device}]")
        choice = input("按Enter使用默认设备，或输入设备ID选择其他设备: ").strip()
        if not choice:
            device_id = default_device
        else:
            try:
                device_id = int(choice)
                if device_id not in input_devices:
                    print(f"错误：无效的设备ID {device_id}")
                    return
            except ValueError:
                print("错误：请输入有效的数字ID")
                return
    else:
        try:
            device_id = int(input("\n请选择麦克风设备ID: "))
            if device_id not in input_devices:
                print(f"错误：无效的设备ID {device_id}")
                return
        except ValueError:
            print("错误：请输入有效的数字ID")
            return
    
    # 开始录音
    recognizer.start_recording(device_id)
    
    # 键盘控制循环
    try:
        while True:
            if WINDOWS:
                # Windows系统使用msvcrt进行键盘检测
                if msvcrt.kbhit():
                    key = msvcrt.getch()
                    
                    # 空格键 - ASCII 32
                    if key == b' ':
                        if recognizer.recording:
                            recognizer.stop_recording()
                        else:
                            recognizer.start_recording(device_id)
                    
                    # Esc键 - ASCII 27
                    elif key == b'\x1b':
                        print("\n退出程序...")
                        if recognizer.recording:
                            recognizer.stop_recording()
                        break
            else:
                # 非Windows系统简化为等待Enter键
                input("\n按Enter键停止录音...")
                if recognizer.recording:
                    recognizer.stop_recording()
                
                choice = input("继续录音?(y/n): ")
                if choice.lower() != 'y':
                    print("\n退出程序...")
                    break
                else:
                    recognizer.start_recording(device_id)
            
            # 短暂睡眠以减少CPU使用
            time.sleep(0.1)
    
    except KeyboardInterrupt:
        print("\n程序被中断")
    finally:
        # 确保资源被释放
        recognizer.close()
    
    print("\n程序已结束。按Enter键退出...")
    input()

if __name__ == "__main__":
    main() 